from queue import Queue
from threading import Thread, Lock
import time
import random

DURATION = 0.2
mutex4print = Lock()


def download(item):
    with mutex4print:
        print(f'Downloading ({item})')
    time.sleep(random.random() * DURATION * 1)
    with mutex4print:
        print(f'Downloaded ({item})')
    return f'Downloaded ({item})'


def resize(item):
    with mutex4print:
        print(f'Resizing ({item})')
    time.sleep(random.random() * DURATION * 3)
    with mutex4print:
        print(f'Resized ({item})')
    return f'Resized ({item})'


def upload(item):
    with mutex4print:
        print(f'Uploading ({item})')
    time.sleep(random.random() * DURATION * 1)
    with mutex4print:
        print(f'Uploaded ({item})')
    return f'Uploaded ({item})'


class ClosableQueue(Queue):

    SENTINEL = object()

    def __init__(self, name):
        super().__init__()
        self.name = name

    def close(self):
        self.put(self.SENTINEL)

    def __iter__(self):
        while True:
            item = self.get()
            try:
                if item is self.SENTINEL:
                    with mutex4print:
                        print(f'<<<< {self.name} encounters SENTINEL.')
                    return
                yield item
            finally:
                with mutex4print:
                    print(f']] Finally {self.name} task_done() for {item}')
                self.task_done()


class StoppableWorker(Thread):

    def __init__(self, func, in_queue, out_queue):
        super().__init__()
        self.func = func
        self.in_queue = in_queue
        self.out_queue = out_queue

    def run(self):
        for item in self.in_queue:
            result = self.func(item)
            self.out_queue.put(result)


if '__main__' == __name__:

    download_queue = ClosableQueue('DOWNLOAD')
    resize_queue = ClosableQueue('RESIZE')
    upload_queue = ClosableQueue('UPLOAD')
    done_queue = ClosableQueue('DONE')
    threads = [
        StoppableWorker(download, download_queue, resize_queue),
        StoppableWorker(resize, resize_queue, upload_queue),
        StoppableWorker(upload, upload_queue, done_queue),
    ]

    for thread in threads:
        thread.start()
    N = 4
    for i in range(N):
        n = i + 1
        download_queue.put(f'obj-{n:04d}')

    download_queue.close()  # 添加N个数据的过程在上面已经结束，不管这个queue空没空，都是加入到最后
    with mutex4print:
        print(f'>>>> Put SENTINEL into downloading queue')
    download_queue.join()
    with mutex4print:
        print(f']]]] Joined downloading queue.')
    resize_queue.close()  # 这里已经通过download_queue.join()来确保download_queue已经结束。
    with mutex4print:
        print(f'>>>> Put SENTINEL into resizing queue')
    resize_queue.join()
    with mutex4print:
        print(f']]]] Joined resizing queue.')
    upload_queue.close()
    with mutex4print:
        print(f'>>>> Put SENTINEL into uploading queue')
    upload_queue.join()
    with mutex4print:
        print(f']]]] Joined uploading queue.')
    with mutex4print:
        print(f'{done_queue.qsize()} items uploaded.')

    for thread in threads:
        thread.join()

    with mutex4print:
        print('All over')
